/*
 * @Description: 
 * @Author: Ren Qian
 * @Date: 2020-03-01 18:07:42
 */

#include "graph_optimizer/optimizer/optimizer.hpp"

namespace graph_optimizer
{
    Optimizer::Optimizer(const std::string &solver_type)
    {
        graph_ptr_.reset(new g2o::SparseOptimizer());

        g2o::OptimizationAlgorithmFactory *solver_factory = g2o::OptimizationAlgorithmFactory::instance();
        g2o::OptimizationAlgorithmProperty solver_property;
        g2o::OptimizationAlgorithm *solver = solver_factory->construct(solver_type, solver_property);
        graph_ptr_->setAlgorithm(solver);

        if (!graph_ptr_->solver())
        {
            std::cerr << "G2O 优化器创建失败！";
        }
        robust_kernel_factory_ = g2o::RobustKernelFactory::instance();
        stamps.clear();
        robust_kernel_name_ = "NONE";
        need_robust_kernel_ = false;
    }

    bool Optimizer::Optimize()
    {
        static int optimize_cnt = 0;
        if (graph_ptr_->edges().size() < 1)
        {
            return false;
        }

        graph_ptr_->initializeOptimization();
        graph_ptr_->computeInitialGuess();
        graph_ptr_->computeActiveErrors();
        graph_ptr_->setVerbose(false);

        double chi2 = graph_ptr_->chi2();
        int iterations = graph_ptr_->optimize(max_iterations_num_);

        std::cout << std::endl
                  << "------ 完成第 " << ++optimize_cnt << " 次后端优化 -------" << std::endl
                  << "顶点数：" << graph_ptr_->vertices().size() << ", 边数： " << graph_ptr_->edges().size() << std::endl
                  << "迭代次数： " << iterations << "/" << max_iterations_num_ << std::endl
                  << "优化前后误差变化：" << chi2 << "--->" << graph_ptr_->chi2()
                  << std::endl
                  << std::endl;

        return true;
    }

    bool Optimizer::GetOptimizedPose(std::deque<PoseData> &optimized_pose)
    {
        optimized_pose.clear();
        int vertex_num = graph_ptr_->vertices().size() / 3;

        for (int i = 0; i < vertex_num; i++)
        {
            PoseData pose;
            pose.time = stamps[i];
            pose.state.p = (dynamic_cast<g2o::VertexP *>(graph_ptr_->vertex(3 * i)))->estimate();
            pose.state.q = (dynamic_cast<g2o::VertexQ *>(graph_ptr_->vertex(3 * i + 1)))->estimate();
            pose.state.bg = (dynamic_cast<g2o::VertexBg *>(graph_ptr_->vertex(3 * i + 2)))->estimate();
            optimized_pose.push_back(pose);
        }
        return true;
    }

    int Optimizer::GetNodeNum()
    {
        return graph_ptr_->vertices().size();
    }

    void Optimizer::SetEdgeRobustKernel(std::string robust_kernel_name,
                                        double robust_kernel_size)
    {
        robust_kernel_name_ = robust_kernel_name;
        robust_kernel_size_ = robust_kernel_size;
        need_robust_kernel_ = true;
    }

    void Optimizer::AddNode(const PoseData &pose)
    {
        stamps.push_back(pose.time);
        g2o::VertexP *vp(new g2o::VertexP());
        vp->setId(graph_ptr_->vertices().size());
        vp->setEstimate(pose.state.p);
        graph_ptr_->addVertex(vp);

        g2o::VertexQ *vq(new g2o::VertexQ());
        vq->setId(graph_ptr_->vertices().size());
        vq->setEstimate(pose.state.q);
        graph_ptr_->addVertex(vq);

        g2o::VertexBg *vbg(new g2o::VertexBg());
        vbg->setId(graph_ptr_->vertices().size());
        vbg->setEstimate(pose.state.bg);
        graph_ptr_->addVertex(vbg);
    }

    void Optimizer::AddOgEdge(int vertex_index1,
                              int vertex_index2,
                              const std::shared_ptr<PreIntegration> &pre_integration)
    {
        g2o::VertexP *vp1 = dynamic_cast<g2o::VertexP *>(graph_ptr_->vertex(vertex_index1));
        g2o::VertexQ *vq1 = dynamic_cast<g2o::VertexQ *>(graph_ptr_->vertex(vertex_index1 + 1));
        g2o::VertexBg *vbg1 = dynamic_cast<g2o::VertexBg *>(graph_ptr_->vertex(vertex_index1 + 2));

        g2o::VertexP *vp2 = dynamic_cast<g2o::VertexP *>(graph_ptr_->vertex(vertex_index2));
        g2o::VertexQ *vq2 = dynamic_cast<g2o::VertexQ *>(graph_ptr_->vertex(vertex_index2 + 1));
        g2o::VertexBg *vbg2 = dynamic_cast<g2o::VertexBg *>(graph_ptr_->vertex(vertex_index2 + 2));

        g2o::EdgeOg *edge(new g2o::EdgeOg());
        edge->setMeasurement(pre_integration);
        // edge->vertices().push_back(vp1);
        // edge->vertices().push_back(vq1);
        // edge->vertices().push_back(vbg1);
        // edge->vertices().push_back(vp2);
        // edge->vertices().push_back(vq2);
        // edge->vertices().push_back(vbg2);
        edge->vertices()[0] = vp1;
        edge->vertices()[1] = vq1;
        edge->vertices()[2] = vbg1;
        edge->vertices()[3] = vp2;
        edge->vertices()[4] = vq2;
        edge->vertices()[5] = vbg2;
        graph_ptr_->addEdge(edge);
        if (need_robust_kernel_)
        {
            AddRobustKernel(edge, robust_kernel_name_, robust_kernel_size_);
        }
    }

    void Optimizer::AddGnssEdge(int gnss_vertex_index,
                                const Eigen::Vector3d &xyz)
    {
        Eigen::Matrix3d information_matrix = Eigen::Matrix3d::Identity() / (gps_noise * gps_noise);
        g2o::VertexP *vp = dynamic_cast<g2o::VertexP *>(graph_ptr_->vertex(gnss_vertex_index));
        g2o::EdgeGnss *edge(new g2o::EdgeGnss());
        edge->setMeasurement(xyz);
        edge->setInformation(information_matrix);
        edge->vertices()[0] = vp;
        graph_ptr_->addEdge(edge);
    }

    void Optimizer::SetMaxIterationsNum(int max_iterations_num)
    {
        max_iterations_num_ = max_iterations_num;
    }

    void Optimizer::AddRobustKernel(g2o::OptimizableGraph::Edge *edge, const std::string &kernel_type, double kernel_size)
    {
        if (kernel_type == "NONE")
        {
            return;
        }

        g2o::RobustKernel *kernel = robust_kernel_factory_->construct(kernel_type);
        if (kernel == nullptr)
        {
            std::cerr << "warning : invalid robust kernel type: " << kernel_type << std::endl;
            return;
        }

        kernel->setDelta(kernel_size);
        edge->setRobustKernel(kernel);
    }
} // namespace graph_optimizer
